import numpy as np 
from pypcd import pypcd
import cv2

bev_image_range=[-30, 30, -30, 30, 1]
bev_image_size=[1080,1920]

def point_to_bev(pointcloud_data_np):
    mask_w = (pointcloud_data_np[:, 0] > bev_image_range[0]+0.1) & (pointcloud_data_np[:, 0] < bev_image_range[1]-0.1)
    mask_h = (pointcloud_data_np[:, 1] > bev_image_range[2]+0.1) & (pointcloud_data_np[:, 1] < bev_image_range[3]-0.1)
    mask_z = pointcloud_data_np[:, 2] < bev_image_range[4]
    mask = mask_w & mask_h & mask_z
    pointcloud_data_np = pointcloud_data_np[mask]

    delta_x_pix_pre_meter = bev_image_size[0] / abs(bev_image_range[0]-bev_image_range[1])
    delta_y_pix_pre_meter = bev_image_size[1] / abs(bev_image_range[2]-bev_image_range[3])
    # 通过激光点的x,y计算当前点在bev上的u,v信息
    back_camera_range_rate=abs(bev_image_range[1]/(bev_image_range[0]-bev_image_range[1]))
    left_camera_range_rate=abs(bev_image_range[3]/(bev_image_range[2]-bev_image_range[3]))
    pointcloud_data_np[:, 0] = np.floor(pointcloud_data_np[:, 0] * delta_x_pix_pre_meter) + bev_image_size[0]*back_camera_range_rate
    pointcloud_data_np[:, 1] = -np.floor(pointcloud_data_np[:, 1] * delta_y_pix_pre_meter) + bev_image_size[1]*left_camera_range_rate
    
    # 将intensity转换为rgb格式
    if np.max(pointcloud_data_np[:,3])<=1.0:
        pointcloud_data_np[:,3]=pointcloud_data_np[:,3]*255.0
    mask_intensity_1 = (pointcloud_data_np[:,3] >= 0) & (pointcloud_data_np[:,3] < 15)
    
    mask_intensity_2 = (pointcloud_data_np[:,3] >= 15) & (pointcloud_data_np[:,3] < 90)
    mask_intensity_3 = (pointcloud_data_np[:,3] >= 90) & (pointcloud_data_np[:,3] < 150)
    mask_intensity_4 = pointcloud_data_np[:,3] >= 150
    color_1_r = 0x00 * mask_intensity_1
    color_1_g = (pointcloud_data_np[:,3] * 255 / 15) * mask_intensity_1
    color_1_b = 0xff * mask_intensity_1

    color_2_r = 0x00 * mask_intensity_2
    color_2_g = 0xff * mask_intensity_2
    color_2_b = ((90 - pointcloud_data_np[:,3]) * 255 / 75) * mask_intensity_2

    color_3_r = ((pointcloud_data_np[:,3] - 90) * 255 / 60) * mask_intensity_3
    color_3_g = 0xff * mask_intensity_3
    color_3_b = 0x00 * mask_intensity_3

    color_4_r = 0x00 * mask_intensity_4
    color_4_g = ((255 - pointcloud_data_np[:,3]) * 255 / (255 - 150)) * mask_intensity_4
    color_4_b = 0xff * mask_intensity_4

    # 每个点的rgb颜色
    cloud_color = np.zeros([len(pointcloud_data_np), 3], np.uint8)
    cloud_color[:,2] = color_1_r + color_2_r + color_3_r + color_4_r
    cloud_color[:,1] = color_1_g + color_2_g + color_3_g + color_4_g
    cloud_color[:,0] = color_1_b + color_2_b + color_3_b + color_4_b

    # 创建以0填充的点云图像
    bev_pointcloud_image_np = np.zeros([bev_image_size[1], bev_image_size[0], 3], np.uint8)
    # 获取点云图像的所有u和v
    u = pointcloud_data_np[:, 0].astype(np.int32)
    v = pointcloud_data_np[:, 1].astype(np.int32)
    # 点云图像中的u，v的点使用对应的颜色填充
    bev_pointcloud_image_np[v, u] = cloud_color.astype(np.uint8)
    return bev_pointcloud_image_np


if __name__=="__main__":
    pc=pypcd.PointCloud.from_path("/home/xd/data/rcs_data/samples/20240902/test/lidar_pcd/171948553700001.pcd")
    xyzi = pc.pc_data[['x', 'y', 'z','intensity']]
    xyzi = np.empty((pc.points, 4), dtype=np.float32)
    xyzi[:, 0] = pc.pc_data['x']
    xyzi[:, 1] = pc.pc_data['y']
    xyzi[:, 2] = pc.pc_data['z']
    xyzi[:, 3] = pc.pc_data['intensity']
    bev_image=point_to_bev(xyzi)
    cv2.imwrite("test.jpg",bev_image)
    print(np.array(xyzi))